from astro.cosmology import Cosmology,PC
from astro.utilities import ismember, indexnear
from astro.io import readtxt
from astro.coord import ang_sep, match
import astro.spec
import gas_gal_hiz.read
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields

wlya = 1215.6701
wlyb = 1025.72
Ckms = 3.e5

GLOBAL_TRANS = 0.76

# plotting limits for transmissivity
YMIN, YMAX = -0.05, 1.2
XMIN, XMAX = -0.5,  15.0
A4LANDSCAPE = 11.7, 8.3

# maximum separation allowed for an LBG from a QSO sight-line
# (degrees).
MAXSEP = 15. / 60.

cosmo = Cosmology(H0=100., Om=0.3, Ol=0.7)

# bins to use when measuring mean flux. distances are in Mpc
maxdist = 20.0
binwidth = 0.5
binedges = np.arange(0, maxdist + binwidth, binwidth)
cbins = binedges[:-1] + 0.5*binwidth
binedges = list(binedges)

# id numbers of qsos that were discovered with VIMOS (so are too faint
# to have aaomega spectra)
num_vimos = np.array([19, 102, 103, 123, 124, 125, 126, 127,
                         143, 144, 145, 146])

# id numbers of qsos observed with UVES/HIRES
num_hires = np.array([5, 16, 60, 101, 131, 134])

def read_data():
    """ Read in qsos, lbgs and qso spectra
    """
    # read LBG redshifts, ras and decs, find LBG comoving LOS
    # distances. (H0=100)
    #lbgs = gas_gal_hiz.read.lbgs('all')

    # now only keep lbgs that are in rich's other list
    lbgs = readtxt(
        '/home/nhmc/projects/gas_gal_hiz/lbgs/LBGall_0.5_data_radecz.dat',
        names='num,ra,dec,z')
    #imatch = match(lbgs.ra,lbgs.dec, richlbgs.ra, richlbgs.dec, 2.0)
    const = 1. / (PC * 1.e6)
    losdist = [cosmo.Dc(z) * const  for z in lbgs.z]      # Mpc
    lbgs = rec_append_fields(lbgs, 'losdist', losdist)

    # read QSO ra, dec and emission redshift (note min(zqso) = 2.25)
    qsos0 = gas_gal_hiz.read.qsos('all')

    # find subset of QSOs for cross correlation - all that are within
    # a few arcmin of an LBG.
    c0 = []
    for q in qsos0:
        # separation is in degrees
        sep = ang_sep(q['ra'], q['dec'], lbgs.ra, lbgs.dec)
        minsep = sep.min()
        c0.append(True if minsep < MAXSEP else False)

    c0 = np.array(c0)
    # remove all that only have VIMOS spectra
    c1 = ~ismember(qsos0.num, num_vimos)
    
    qsos1 = qsos0[c0 & c1]

    # take a look
    #plot(lbgs.ra,lbgs.dec,'g.')
    #plot(qsos1.ra,qsos1.dec,'o', mfc='None')
    #pl.axis('equal')

    # now find high resolution spectra
    is_hires_spec = ismember(qsos1.num, num_hires)
    qsos_hires = qsos1[is_hires_spec]
    # and aaomega spectra
    qsos_aaom = qsos1[~is_hires_spec]
    print len(qsos_hires), len(qsos_aaom)

    # read the high res spectra
    hires_filenames = {
        5   : '/home/nhmc/projects/gas_gal_hiz/qsos/Q0042m2627.dat',
        16  : '/home/nhmc/projects/gas_gal_hiz/qsos/Q0043m265cont.dat',
        60  : '/home/nhmc/projects/gas_gal_hiz/qsos/J0124p0044.asc',
        101 : '/home/nhmc/projects/gas_gal_hiz/qsos/HE0940m1050m.asc',
        131 : '/home/nhmc/projects/gas_gal_hiz/qsos/Q2129m1602.asc',
        134 : '/home/nhmc/projects/gas_gal_hiz/qsos/PKS2126m158.asc'}

    spec_hires = []
    for qso in qsos_hires:
        filename = hires_filenames[qso['num']]
        print 'reading %s' % filename
        spec_hires.append(astro.spec.read(filename))

    # read the aaomega spectra
    prefix = '/home/nhmc/projects/gas_gal_hiz/qsos/aaomega/'
    aaomdirs = dict(Q0042   = prefix + 'q0042/',
                    J0124   = prefix + 'j0124/',
                    HE0940  = prefix + 'he0940/',
                    PKS2126 = prefix + 'pks2126/',
                    J1201   = prefix + 'j1201/')

    spec_aaom = []
    for qso in qsos_aaom:
        filename = aaomdirs[qso['fieldqso']] + qso['name'] + '.txt'
        print 'reading %s' % filename
        spec_aaom.append(astro.spec.read(filename))
    
    return lbgs, qsos_hires, qsos_aaom, spec_hires, spec_aaom

def calc_trans(binedges,
               lbgs_angdist, lbgs_losdist, lbgs_num,
               lyalosdist, nfl, masked):
    """ Given a list of lbgs, LOS distances, normalised fluxes, masks
    and bins, returns mean flux per bin, mean flux per galaxy per bin,
    and the galaxy id numbers for each bin.
    """
    nbins = len(binedges) - 1
    
    # mean norm flux per galaxy per bin
    meanfl_gal_bins = [[] for i in range(nbins)] 
    # the id numbers of gals that contribute to each bin (for later checks)
    idgal = [[] for i in range(nbins)]
    # loop over all lbgs and bins.
    lenlbgs = len(lbgs_angdist)
    for j,(angdist, losdist) in enumerate(zip(lbgs_angdist,lbgs_losdist)):
        #print '%s of %s lbgs' % (j+1, lenlbgs)
        print '.',
        # separations between lbg position and all QSO pixel positions
        sep = np.hypot(angdist, lyalosdist - losdist)
        for i in xrange(nbins):
            # find all pixel values that fall in this distance bin
            # that are not masked
            c0 = (binedges[i] <= sep) & (sep < binedges[i+1])
            nfl1 = nfl[c0 & ~masked]
            npix = len(nfl1)
            if npix > 0:
                #print 'contributes to bin %i' % i 
                nfl_mean = nfl1.sum() / npix
                meanfl_gal_bins[i].append(nfl_mean)
                idgal[i].append(lbgs_num[j])

    # no. of LBGs contributing each bin
    return meanfl_gal_bins, idgal

def trans_setup(qso, spec, lbgs, zscale=None):
    """ Select appropiate QSO wav range and LBGs overlapping that
    range, Calculate comoving separation of LBGs from QSO sightline,
    and LOS comoving distance of the QSO pixels.

    Make a boolean mask to flag bad flux values.

    Scale the normalised flux if required.
    """
    # find minimum and maximum redshifts based on Lya forest region:
    dv = 3000.
    wmin = wlyb * (1.+ qso['z'])
    wmax = wlya * (1 + qso['z']) * (1 - dv/Ckms)

    dv = 5000.
    # go a little further beyond the z limits to look for LBGs.
    zmin = (wmin / wlya) * (1 - dv/Ckms) - 1
    zmax = (wmax / wlya) * (1 + dv/Ckms) - 1
    print 'Redshift limits %5.2f -> %5.2f' % (zmin,zmax)

    # lbg selection
    c0 = (zmin <= lbgs['z']) & (lbgs['z'] <= zmax)
    lbgs0 = lbgs[c0]
    sep = ang_sep(qso['ra'], qso['dec'], lbgs0.ra, lbgs0.dec)
    is_close =  sep < MAXSEP
    lbgs1 = lbgs0[is_close]
    # angular separation of lbgs from QSO sightline (radians)
    lbg_angsep = sep[is_close] * np.pi / 180.
    # comoving transverse separation 
    lbg_angdist = np.array([r*los for r,los in zip(lbg_angsep, lbgs1.losdist)])

    # only keep lbgs with separations that are smaller than our
    # largest separation bin limits.
    #c0 = lbg_angdist < max(binedges)
    #lbg_angdist = lbg_angdist[c0]
    #lbgs0 = lbgs0[c0]

    #if len(lbgs0) == 0:
    #    return (None,) * 5

    # select region of QSO spectrum we're interested in and normalise.
    spec0 = spec[(wmin <= spec.wa) & (spec.wa <= wmax)] 
    wa = spec0.wa
    nfl0 = spec0.fl / spec0.co

    ner = spec0.er / spec0.co
    zlya = wa / wlya - 1.
    const = 1 / (PC * 1.e6)
    lyalosdist = np.array([cosmo.Dc(z) * const for z in zlya]) 
    if zscale == False:
        nfl1 = nfl0  
    elif zscale == True:
        scaleval = GLOBAL_TRANS / (0.676 - 0.220*(zlya-3.))
        print 'Mean flux scaling value %.3f' % scaleval.mean()
        nfl1 = nfl0 * scaleval
    elif zscale == 'mean':
        good = (ner > 0) & ~np.isnan(nfl0)
        scaleval = GLOBAL_TRANS / nfl0[good].mean()
        print 'Flux scaling value %.3f' % scaleval
        nfl1 = nfl0 * scaleval
    else:
        raise ValueError('Unknown zscale value: %s' % zscale)
    
    # masking regions. where ma=0 pixel is masked.
    masked = np.ones(len(wa), dtype=bool)
    #if masks:
    #   for wmin,wmax in masks:
    #       print 'masking from %s to %s' % (wmin,wmax)
    #       condition = (wmin < wa) & ( wa < wmax)
    #       masked = np.where(condition, True, masked)
    # unmask good points
    masked[(ner > 0) & ~np.isnan(nfl1)] = False
    print '%i of %i  pixels masked' % (len(masked.nonzero()[0]), len(masked))

    return lyalosdist, nfl1, masked, lbgs1, lbg_angdist


def run_trans_and_plot(qsos, spectra, lbgs, title=None, zscale=None,
                       per_spec_plots=False, singleplot=False,
                       info='hires'):
    """ Calculate the mean flux near LBGs for a list of QSOs, and so
    some plots.
    """
    meanfl_per_gal_per_bin = [[] for i in range(len(cbins))]
    idgal_per_bin = [[] for i in range(len(cbins))]

    cols = 'brgycm'
    if zscale == 'mean': note = 'scalemean'
    elif zscale == True: note = 'scalefz'
    else: note = 'noscaling'

    if singleplot:
        fig = pl.figure(figsize=A4LANDSCAPE)
        ax = fig.add_subplot(111)
        ax.axhline(GLOBAL_TRANS, color='gray')

    for i,(qso,spec) in enumerate(zip(qsos, spectra)):
        print '\nCalculating trans for %(name)s' % qso 
        lyalosdist, nfl, masked, lbgs0, lbg_angdist = trans_setup(
            qso, spec, lbgs, zscale=zscale)

        if lyalosdist is None:
            print 'There are no LBGs close enough to the forest, skipping'
            continue
        meanfl_per_gal, idgals = calc_trans(
            binedges, lbg_angdist, lbgs0.losdist, lbgs0.num,
            lyalosdist, nfl, masked)
        
        if per_spec_plots:
            temp = '%(name)s_%(num)i' % qso
            filename = '%s_meanfl_%s_%s.pdf' % (info, temp, note)
            do_errplot(meanfl_per_gal, '%s %s' % (note, temp), filename)
        if singleplot:
            ngal = np.array([len(g) for g in idgals])
            sumfl = np.array([sum(f) for f in meanfl_per_gal])
            meanfl = sumfl / ngal
            c = cols[i % len(cols)]
            ax.plot(cbins, meanfl, '-o' + c, label=qso['name'])

        # add the mean flux values and galaxy ids to a large list.
        for j in range(len(cbins)):
            idgal_per_bin[j].extend(idgals[j])
            meanfl_per_gal_per_bin[j].extend(meanfl_per_gal[j])

    if singleplot:
        ax.set_title(title)
        ax.set_ylabel('Normalised Flux')
        ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
        ax.grid()
        ax.minorticks_on()
        ax.set_ylim(YMIN, YMAX)
        ax.set_xlim(XMIN, XMAX)
        if len(spectra) < 10:
            ax.legend(loc='lower right')
        fig.savefig('pdf/meanfl_%s_%s.pdf' % (info, note))        
        fig.savefig('meanfl_%s_%s.png' % (info, note))        
        
    #pl.show()

    return meanfl_per_gal_per_bin, idgal_per_bin


def do_boxplot(meanfl_per_gal, title, filename):
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ax.axhline(GLOBAL_TRANS, color='gray')
    ax.boxplot(meanfl_per_gal, positions=cbins)
    ax.set_title(title)
    ax.set_ylabel('Normalised Flux')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_ylim(YMIN, YMAX)
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))

def do_errplot(meanfl_per_gal, title, filename, show=False):
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ax.axhline(GLOBAL_TRANS, color='gray')
    std =  np.array([np.std(f) for f in meanfl_per_gal])
    nvals = np.array([len(f) for f in meanfl_per_gal])
    stderr = std / np.sqrt(nvals)
    mean = [np.mean(f) for f in meanfl_per_gal]
    ax.errorbar(cbins, mean, yerr=stderr, fmt='o-b')
    ax.set_title(title)
    ax.set_ylabel('Normalised Flux')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_ylim(YMIN, YMAX)
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))
    if not show:
        pl.close()

def do_ngalplot(meanfl_per_gal, title, filename, show=False):
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    ngal = [len(f) for f in meanfl_per_gal]
    pl.semilogy(cbins, ngal, ls='steps-mid')
    ax.set_title(title)
    ax.set_ylabel('Number of galaxy-LOS pairs')
    ax.set_xlabel('Separation / ($h^{-1}$ Mpc)')
    ax.grid()
    ax.minorticks_on()
    ax.set_xlim(XMIN, XMAX)
    fig.savefig('pdf/' + filename)
    fig.savefig(filename.replace('.pdf','.png'))
    if not show:
        pl.close()
